import numpy as np
import random

import torch
from PIL import Image
from torchvision import transforms as T
from torchvision.transforms import functional as F


def pad_if_smaller(img, size, fill=0):
    # 如果图像最小边长小于给定size，则用数值fill进行padding
    min_size = min(img.size)
    if min_size < size:
        ow, oh = img.size
        padh = size - oh if oh < size else 0
        padw = size - ow if ow < size else 0
        img = F.pad(img, (0, 0, padw, padh), fill=fill)
    return img


class Compose(object):
    def __init__(self, transforms):
        self.transforms = transforms

    def __call__(self, image, target):
        for t in self.transforms:
            image, target = t(image, target)
        return image, target


class RandomResize(object):
    def __init__(self, min_size, max_size=None):
        self.min_size = min_size
        if max_size is None:
            max_size = min_size
        self.max_size = max_size

    def __call__(self, image, target):
        size = random.randint(self.min_size, self.max_size)   #size是一个包含两个整数的元组 (width, height)，表示期望的宽度和高度。
        # 这里size传入的是int类型，所以是将图像的最小边长缩放到size大小
        image = F.resize(image, size)
        # 这里的interpolation注意下，在torchvision(0.9.0)以后才有InterpolationMode.NEAREST
        # 如果是之前的版本需要使用PIL.Image.NEAREST
        target = F.resize(target, size, interpolation=T.InterpolationMode.NEAREST)
        return image, target


class RandomHorizontalFlip(object):
    def __init__(self, flip_prob):
        self.flip_prob = flip_prob

    def __call__(self, image, target):
        if random.random() < self.flip_prob:
            image = F.hflip(image)
            target = F.hflip(target)
        return image, target


class RandomVerticalFlip(object):
    def __init__(self, flip_prob):
        self.flip_prob = flip_prob

    def __call__(self, image, target):
        if random.random() < self.flip_prob:
            image = F.vflip(image)
            target = F.vflip(target)
        return image, target

class RandomRotate(object):
    def __init__(self, angles):
        self.angles = angles

    def __call__(self, img, target):
        angle = random.choice(self.angles)
        return img.rotate(angle, resample=Image.BILINEAR), target.rotate(angle, resample=Image.NEAREST)


class RandomCrop(object):
    def __init__(self, size):
        self.size = size

    def __call__(self, image, target):
        image = pad_if_smaller(image, self.size)
        target = pad_if_smaller(target, self.size, fill=255)                               #target填充的数值是没有意义的，所以用255填充
        crop_params = T.RandomCrop.get_params(image, (self.size, self.size))  #用于生成随机裁剪操作所需的参数，包括裁剪框的位置和大小。该方法返回一个包含裁剪参数的字典，其中包括以下键值对：i: 水平方向上裁剪框左上角的 x 坐标。j: 垂直方向上裁剪框左上角的 y 坐标。h: 裁剪框的高度。w: 裁剪框的宽度。
        image = F.crop(image, *crop_params)
        target = F.crop(target, *crop_params)
        return image, target


class CenterCrop(object):
    def __init__(self, size):
        self.size = size

    def __call__(self, image, target):
        image = F.center_crop(image, self.size)
        target = F.center_crop(target, self.size)
        return image, target


class ToTensor(object):
    def __call__(self, image, target):
        image = F.to_tensor(image)
        target = torch.as_tensor(np.array(target), dtype=torch.int64)            #target记录的是每个像素的类别索引，不需要对其进行缩放，所以转化为tensor就行了，torch.as_tensor() 是 PyTorch 提供的一个函数，可以将 Python 数组、NumPy 数组或其他数组类型转换为 PyTorch 张量。
        return image, target


class Normalize(object):
    def __init__(self, mean, std):
        self.mean = mean
        self.std = std

    def __call__(self, image, target):
        image = F.normalize(image, mean=self.mean, std=self.std)
        #img_arr = np.array(image)
        return image, target
